Hmm, looks like punet fails to compile with this. I think it must be one of the conv dispatches (because nothing else uses this pass), but it seems a bit strange to me that this is causing a compilation failure.

Hmm, looks like punet fails to compile with this. I think it must be one of the conv dispatches (because nothing else uses this pass), but it seems a bit strange to me that this is causing a compilation failure.

Yeah thats a bummer, one thing I can try is to see if CI passes when I disable multi-result, if thats the case then the issue might be that the getTiedResult function assumes the multiple parallel_insert_slice are ordered in the same way as the results.

I took a look at the model, and found the dispatch that is failing. Here's a gist with a dump: https://gist.github.com/Max191/a092ec5c81e5f9f44aee19358f007a34

Looks like the problem is that the collapse_shape was hoisted out of the forall, but then could not fold with the flow.dispatch.tensor.store at the function boundary, because the store has an offset (from pad fusion). I'm not quite sure what the best solution for this is at the moment.

EDIT: snippet from the dump after reshape propagation:

    } -> (tensor<1x2x8x16x4x16xf16>, tensor<1x2x8x16x4x16xi8>)
    scf.forall.in_parallel {
      tensor.parallel_insert_slice %27#0 into %arg3[%arg0, %arg1, 0, 0, %arg2, 0] [1, 2, 8, 16, 4, 16] [1, 1, 1, 1, 1, 1] : tensor<1x2x8x16x4x16xf16> into tensor<2x128x8x16x20x16xf16>
      tensor.parallel_insert_slice %27#1 into %arg4[%arg0, %arg1, 0, 0, %arg2, 0] [1, 2, 8, 16, 4, 16] [1, 1, 1, 1, 1, 1] : tensor<1x2x8x16x4x16xi8> into tensor<2x128x8x16x20x16xi8>
    }
  } {mapping = [#iree_codegen.workgroup_mapping<z>, #iree_codegen.workgroup_mapping<y>, #iree_codegen.workgroup_mapping<x>]}
  %collapsed = tensor.collapse_shape %20#1 [[0], [1], [2, 3], [4, 5]] : tensor<2x128x8x16x20x16xi8> into tensor<2x128x128x320xi8>
  flow.dispatch.tensor.store %collapsed, %7, offsets = [0, 1, 1, 0], sizes = [2, 128, 128, 320], strides = [1, 1, 1, 1] : tensor<2x128x128x320xi8> -> !flow.dispatch.tensor<readwrite:tensor<2x130x130x320xi8>>
  flow.dispatch.tensor.store %20#0, %8, offsets = [0, 0, 0, 0, 0, 0], sizes = [2, 128, 8, 16, 20, 16], strides = [1, 1, 1, 1, 1, 1] : tensor<2x128x8x16x20x16xf16> -> !flow.dispatch.tensor<writeonly:tensor<2x128x8x16x20x16xf16>>
  return
}

One way to fix this would be to check that the consumer of the scf.forall is a flow.dispatch.tensor.store with a full slice (i.e., zero offsets, unit strides, and source_sizes == target_sizes). This is making the pattern even more specific, but it is already quite specific, so maybe it is okay for now.

One way to fix this would be to check that the consumer of the scf.forall is a flow.dispatch.tensor.store with a full slice (i.e., zero offsets, unit strides, and source_sizes == target_sizes). This is making the pattern even more specific, but it is already quite specific, so maybe it is okay for now.

Makes sense, I guess theoretically it should be possible to support this?

flow.dispatch.tensor.store %20#1, %7, offsets = [0, 1, 0, 1, 0, 0], sizes = [2, 128, 8, 16, 20, 16], strides = [1, 1, 1, 1, 1, 1] : tensor<2x128x8x16x20x16xf16> -> !flow.dispatch.tensor<writeonly:tensor<2x130x130x20x16xf16>

Following up here from a discussion with Nirvedh offline. We are going with the solution I proposed here: #19044 (comment)

As long as the convolution has a consumer in the dispatch, then the pipeline should be able to handle not hoisting the collapse_shape when the dispatch.tensor.store has offsets (when a pad consumer has been fused). If there is no consumer, and the result of the convolution is directly stored with some offsets, then this will cause problems, but I don't think it is a case that is likely to be seen. There is typically an operation between convolutions and pads, and pad fusion does not happen by default anyway. It is a case to be aware of, but I think it is fine to proceed with this additional restriction for now.

In other words, the following case would be problematic if the pad is fused with the producer convolution:

%conv = linalg.conv_2d_nhwc_hwcf ...
%pad = tensor.pad %conv ...

However, this seems like an uncommon case, and pad fusion only happens with aggressive fusion anyway.